import os
import sys
import argparse
from typing import List

BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
sys.path.append(BASE_DIR)
from utils.sql_app.database import Models_DIR
from conf.config import Config
from aerich import get_app_connection, InspectMySQL, get_app_connection_name
from aerich.inspectdb import Column
from tortoise import connections
from utils.sql_app.database import TORTOISE_ORM
from common import logger
DAL_DIR = os.path.join(Config.BASE_DIR, 'dal', "sql_dal")


class MyColumn(Column):

    def translate(self) -> dict:
        comment = default = length = index = null = pk = ""
        if self.pk:
            pk = "pk=True, "
        else:
            if self.unique:
                index = "unique=True, "
            else:
                if self.index:
                    index = "index=True, "
        if self.data_type in ["varchar", "VARCHAR"]:
            length = f"max_length={self.length}, "
        if self.data_type in ["decimal", "numeric"]:
            length_parts = []
            if self.max_digits:
                length_parts.append(f"max_digits={self.max_digits}, ")
            if str(self.decimal_places):
                length_parts.append(f"decimal_places={self.decimal_places}, ")
            length = "".join(length_parts)
        if self.null:
            null = "null=True, "
        if self.default is not None:
            if self.data_type in ["tinyint", "INT"]:
                default = f"default={self.default}, "
            elif self.data_type == "bool":
                default = f"default={'1' if self.default == 'true' else '0'}, "
            elif self.data_type in ["datetime", "timestamptz", "TIMESTAMP"]:
                if "CURRENT_TIMESTAMP" == self.default:
                    if "DEFAULT_GENERATED on update CURRENT_TIMESTAMP" == self.extra:
                        default = "auto_now=True, "
                    else:
                        default = "auto_now_add=True, "
            else:
                if "::" in self.default:
                    default = f"default={self.default.split('::')[0]}, "
                elif self.default.endswith("()"):
                    default = '""'
                else:
                    self.default = '""' if self.default == "" else self.default
                    default = f"default={self.default}, "

        if self.comment:
            comment = f"description={repr(self.comment)}, "
        return {
            "name": self.name,
            "pk": pk,
            "index": index,
            "null": null,
            "default": default,
            "length": length,
            "comment": comment,
        }


class MyInspectMySQL(InspectMySQL):

    @property
    def field_map(self) -> dict:
        return {
            "int": self.int_field,
            "smallint": self.smallint_field,
            "tinyint": self.int_field,
            "bigint": self.bigint_field,
            "varchar": self.char_field,
            "longtext": self.text_field,
            "text": self.text_field,
            "datetime": self.datetime_field,
            "float": self.float_field,
            "date": self.date_field,
            "time": self.time_field,
            "decimal": self.decimal_field,
            "json": self.json_field,
            "longblob": self.binary_field,
        }

    @classmethod
    def int_field(cls, **kwargs) -> str:
        return "{name} = fields.IntField({pk}{index}{default}{comment})".format(**kwargs)

    async def inspect(self) -> str:
        if not self.tables:
            self.tables = await self.get_all_tables()
        result = "from tortoise import Model, fields\n\n\n"
        tables = []
        for table in self.tables:
            columns = await self.get_columns(table)
            fields, ex_fields = [], []
            model = self._table_template.format(table=table.title().replace("_", ""))
            for column in columns:
                field = self.field_map[column.data_type](**column.translate())
                if field in ex_fields:
                    continue
                ex_fields.append(field)
                fields.append("    " + field.replace(", )", ")"))
            tables.append(model + "\n".join(fields))
        return result + "\n\n\n".join(tables)

    async def get_columns(self, table: str) -> List[Column]:
        columns = []
        sql = """
            select c.*, s.NON_UNIQUE, s.INDEX_NAME
            from information_schema.COLUMNS c
            left join information_schema.STATISTICS s on c.TABLE_NAME = s.TABLE_NAME
            and c.TABLE_SCHEMA = s.TABLE_SCHEMA
            and c.COLUMN_NAME = s.COLUMN_NAME
            where c.TABLE_SCHEMA = %s
            and c.TABLE_NAME = %s
        """
        ret = await self.conn.execute_query_dict(sql, [self.database, table])
        for row in ret:
            non_unique = row["NON_UNIQUE"]
            if non_unique is None:
                unique = False
            else:
                unique = not non_unique
            index_name = row["INDEX_NAME"]
            if index_name is None:
                index = False
            else:
                index = row["INDEX_NAME"] != "PRIMARY"
            columns.append(
                MyColumn(
                    name=row["COLUMN_NAME"],
                    data_type=row["DATA_TYPE"],
                    null=row["IS_NULLABLE"] == "YES",
                    default=row["COLUMN_DEFAULT"],
                    pk=row["COLUMN_KEY"] == "PRI",
                    comment=row["COLUMN_COMMENT"],
                    unique=row["COLUMN_KEY"] == "UNI",
                    extra=row["EXTRA"],
                    unque=unique,
                    index=index,
                    length=row["CHARACTER_MAXIMUM_LENGTH"],
                    max_digits=row["NUMERIC_PRECISION"],
                    decimal_places=row["NUMERIC_SCALE"],
                )
            )
        return columns


async def inspectdb(tables: List[str] = None) -> str:
    await connections._init(TORTOISE_ORM["connections"], False)
    connection = get_app_connection(TORTOISE_ORM, "models")
    inspect = MyInspectMySQL(connection, tables)
    return await inspect.inspect()


if __name__ == '__main__':
    import asyncio

    loop = asyncio.get_event_loop()
    parser = argparse.ArgumentParser(description='generation model')
    parser.add_argument('--table', '-t', type=str, default=None, required=True, help="表名称")

    args = parser.parse_args()
    result = loop.run_until_complete(inspectdb([args.table]))

    result += f"""
    
    class Meta:
        table = '{args.table}'
    """
    with open(f"{Models_DIR}/{args.table}.py", "w+", encoding="utf-8") as f:
        f.write(result)

    logger.info(f"生成 {args.table} 成功")

    if os.path.exists(f"{DAL_DIR}/{args.table}_dal.py"):
        logger.info(f"{args.table}_dal.py 已经存在,不进行覆盖")
        exit(1)

    with open(f"{DAL_DIR}/{args.table}_dal.py", "w+", encoding="utf-8") as f:
        parts = args.table.split('_')
        capitalized_parts = [part.capitalize() if part else part for part in parts]
        '_'.join(capitalized_parts)

        text = f"""
from dal.sql_dal import BaseDal
from dal.sql_dal.models.{args.table} import {''.join(capitalized_parts)}


class {''.join(capitalized_parts)}Dal(BaseDal):
    def __init__(self):
        super().__init__(model={''.join(capitalized_parts)})


# {args.table}_dal = {''.join(capitalized_parts)}Dal()
        """
        f.write(text)